#!/bin/bash

# 颜色和样式设置
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
BOLD='\033[1m'
RESET='\033[0m'

# 函数：创建目录（如果不存在的话）
create_directory() {
    local dir="$1"
    if [ ! -d "$dir" ]; then
        echo -e "${YELLOW}${BOLD}目录 $dir 不存在. 正在创建...${RESET}"
        mkdir -p "$dir"
    fi
}

# 函数：将字符 ':' 和 '/' 替换为 '_'，以便创建有效的文件名
sanitize_filename() {
    echo "$1" | sed 's/:/_/g; s/\//_/g'
}

# 函数：打包所有镜像到指定目录
package_all_images() {
    local dest_dir="$1"
    create_directory "$dest_dir"
    
    echo -e "${BLUE}${BOLD}正在列出所有镜像...${RESET}"
    images=$(podman images --format "{{.Repository}}:{{.Tag}}")
    
    for image in $images; do
        local filename=$(sanitize_filename "$image")
        local filepath="$dest_dir/$filename.tar"
        
        echo -e "${GREEN}${BOLD}正在保存 $image 到 $filepath...${RESET}"
        podman save -o "$filepath" "$image"
        if [ $? -eq 0 ]; then
            echo -e "${GREEN}${BOLD}$image 已成功保存为 $filepath${RESET}"
        else
            echo -e "${RED}${BOLD}$image 保存失败${RESET}"
        fi
    done
    
    echo -e "${GREEN}${BOLD}所有镜像已打包到 $dest_dir。${RESET}"
}

# 函数：打包指定镜像到指定目录
package_specific_image() {
    local image="$1"
    local dest_dir="$2"
    
    create_directory "$dest_dir"
    
    local filename=$(sanitize_filename "$image")
    local filepath="$dest_dir/$filename.tar"
    
    echo -e "${GREEN}${BOLD}正在保存 $image 到 $filepath...${RESET}"
    podman save -o "$filepath" "$image"
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}${BOLD}镜像 $image 已打包到 $dest_dir。${RESET}"
    else
        echo -e "${RED}${BOLD}镜像 $image 打包失败${RESET}"
    fi
}

# 函数：从指定目录导入镜像
import_images() {
    local src_dir="$1"
    
    if [ ! -d "$src_dir" ]; then
        echo -e "${RED}${BOLD}目录 $src_dir 不存在. 退出...${RESET}"
        exit 1
    fi
    
    echo -e "${BLUE}${BOLD}正在从 $src_dir 导入镜像...${RESET}"
    for tar_file in "$src_dir"/*.tar; do
        if [ -f "$tar_file" ]; then
            echo -e "${GREEN}${BOLD}正在导入 $tar_file...${RESET}"
            podman load -i "$tar_file"
            if [ $? -eq 0 ]; then
                echo -e "${GREEN}${BOLD}镜像 $tar_file 成功导入${RESET}"
            else
                echo -e "${RED}${BOLD}镜像 $tar_file 导入失败${RESET}"
            fi
        else
            echo -e "${YELLOW}${BOLD}文件 $tar_file 不存在${RESET}"
        fi
    done
    
    echo -e "${GREEN}${BOLD}所有镜像已从 $src_dir 导入。${RESET}"
}

# 主脚本逻辑
echo -e "${BLUE}${BOLD}请选择一个选项:${RESET}"
echo -e "${BLUE}1. 将所有镜像打包到指定目录${RESET}"
echo -e "${BLUE}2. 将指定镜像打包到指定目录${RESET}"
echo -e "${BLUE}3. 从指定目录导入镜像${RESET}"
read -r option

case $option in
    1)
        echo -e "${BLUE}请输入目标目录:${RESET}"
        read -r dest_dir
        package_all_images "$dest_dir"
        ;;
    2)
        echo -e "${BLUE}请输入镜像名称（例如 'ubuntu:latest'）：${RESET}"
        read -r image_name
        echo -e "${BLUE}请输入目标目录:${RESET}"
        read -r dest_dir
        package_specific_image "$image_name" "$dest_dir"
        ;;
    3)
        echo -e "${BLUE}请输入包含 tar 文件的源目录:${RESET}"
        read -r src_dir
        import_images "$src_dir"
        ;;
    *)
        echo -e "${RED}${BOLD}无效的选项. 退出...${RESET}"
        exit 1
        ;;
esac